import asyncio
import os
import json
import sys
from typing import Optional, Dict, Any
from contextlib import AsyncExitStack

import aiohttp
from openai import OpenAI # type: ignore
from dotenv import load_dotenv # type: ignore
from mcp import ClientSession, StdioServerParameters # type: ignore
from mcp.client.stdio import stdio_client # type: ignore

# 加载 .env 文件,确保 API_KEY 和 API_SECRET 可用
load_dotenv()

class MCPClient:
    def __init__(self, server_url: str = "http://localhost:8080"):
        """初始化 MCP 客户端"""
        self.exit_stack = AsyncExitStack()
        self.openai_api_key = os.getenv("OPENAI_API_KEY")
        self.base_url = os.getenv("BASE_URL")
        self.model = os.getenv("MODEL")
        self.session: Optional[ClientSession] = None
        self.server_url = server_url
        
        # 验证所有必需的环境变量
        if not all([self.openai_api_key, self.base_url, self.model]):
            raise ValueError("❌ 请确保设置了所有必需的环境变量: OPENAI_API_KEY, BASE_URL, MODEL")

        # 确保 base_url 以 '/' 结尾
        if self.base_url and not self.base_url.endswith('/'):
            self.base_url += '/'

        self.client = OpenAI(
            api_key=self.openai_api_key,
            base_url=self.base_url
        )

        self.history = [
            {"role": "system", "content": os.getenv("ROLE")}
        ]
        self.temperature = float(os.getenv("TEMPERATURE", 0.7))

    async def connect_to_server(self, server_script_path: str):
        """连接到 MCP 服务器并列出可用工具"""
        is_python = server_script_path.endswith('.py')
        is_js = server_script_path.endswith('.js')
        if not (is_python or is_js):
            raise ValueError("服务器脚本必须是 .py 或 .js 文件")

        command = "python" if is_python else "node"
        server_params = StdioServerParameters(
            command=command,
            args=[server_script_path],
            env=None
        )

        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.writer = stdio_transport
        self.session = await self.exit_stack.enter_async_context(
            ClientSession(self.stdio, self.writer)
        )

        await self.session.initialize()
        response = await self.session.list_tools()
        print("\n已连接到服务器，可用工具：", [tool.name for tool in response.tools])

    async def process_query(self, query: str) -> str:
        """处理用户查询并支持工具调用"""
        try:
            if not query.strip():
                return "❌ 请输入有效的查询内容"

            # 获取可用工具
            tools_response = await self.session.list_tools()
            available_tools = [{
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.inputSchema
                }
            } for tool in tools_response.tools]

            # 添加用户输入到历史记录
            messages = [{"role": "user", "content": query}]
            
            try:
                # 调用 API 进行工具选择
                response = await asyncio.get_event_loop().run_in_executor(
                    None, 
                    lambda: self.client.chat.completions.create(
                        model=self.model,
                        messages=messages,
                        tools=available_tools,
                        temperature=self.temperature
                    )
                )

                content = response.choices[0]
                if content.finish_reason == "tool_calls":
                    # 处理工具调用
                    tool_call = content.message.tool_calls[0]
                    tool_name = tool_call.function.name
                    tool_args = json.loads(tool_call.function.arguments)
                    
                    print(f"\n[调用工具 {tool_name}，参数：{tool_args}]")
                    result = await self.session.call_tool(tool_name, tool_args)
                    
                    # 将工具调用结果添加到对话
                    messages.append(content.message.model_dump())
                    messages.append({
                        "role": "tool",
                        "content": result.content,
                        "tool_call_id": tool_call.id
                    })
                    
                    # 最终处理
                    final_response = await asyncio.get_event_loop().run_in_executor(
                        None,
                        lambda: self.client.chat.completions.create(
                            model=self.model,
                            messages=messages,
                            temperature=self.temperature
                        )
                    )
                    return final_response.choices[0].message.content
                
                return content.message.content

            except Exception as e:
                error_msg = str(e)
                if "404" in error_msg:
                    return f"❌ API 端点无法访问。请检查 BASE_URL 设置是否正确: {self.base_url}"
                return f"❌ API 调用错误: {error_msg}"
                
        except Exception as e:
            return f"❌ 处理查询时出错: {str(e)}"

    async def chat_loop(self):
        """运行交互式聊天循环"""
        print("\n🚀 MCP客户端已启动！输入'quit'退出")
        print("示例命令：")
        print("1. 查询天气：'查询北京市的天气'")
        print("2. 上传代码：'将代码上传到Gitee仓库'")

        while True:
            try:
                query = input("\n你: ").strip()
                if query == "quit":
                    break
                response = await self.process_query(query)
                print(f"\n[bot]: {response}")
            except Exception as e:
                print(f"❌ 出现错误: {e}")

    async def cleanup(self):
        """清理资源"""
        await self.exit_stack.aclose()

    async def upload_to_gitee(self, repo_path: str, commit_msg: str = None) -> Dict[str, Any]:
        """
        通过MCP服务上传代码到Gitee
        Args:
            repo_path: 本地仓库路径
            commit_msg: 提交信息(可选)
        """
        async with aiohttp.ClientSession() as session:
            async with session.post(
                f"{self.server_url}/api/git/upload",
                json={
                    'repo_path': repo_path,
                    'commit_msg': commit_msg
                }
            ) as resp:
                return await resp.json()

async def main():
    if len(sys.argv) < 2:
        print("用法：python client.py <服务器脚本路径>")
        sys.exit(1)
        
    client = MCPClient()
    try:
        await client.connect_to_server(sys.argv[1])
        await client.chat_loop()
    finally:
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())
